/************************************************************************
Copyright (c) 2010, 2013, C.R. All rights reserved. 

E-mail: crtoyou@gmail.com; crtoyou@163.com

You can redistribute it and/or modify it under the terms of the
GNU Lesser General Public License as published by the
Free Software Foundation version 3 of the License.

This source code is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. if not, write to the Free Software Foundation, Inc., 
51 Franklin St, Fifth Floor, Boston, MA 02110-1301  USA
************************************************************************/

#ifndef __CR_HUFFMAN_TREE_H__
#define __CR_HUFFMAN_TREE_H__

#include <vector>
#include <crbsplus.h>
#include <exception>

// for debug probe
#include <fstream>
#include <iomanip>

#define _SKIP_NO_FREQ_

namespace crhuffman {

using namespace cr;

//////////////////////////////////////////////////////////////////////////
// Huffman tree: 0:left  1:right
// root take no bit in code
//////////////////////////////////////////////////////////////////////////

template <size_t CODE_SPACE /*256|16*/>
class huffman_tree
{
private:
	huffman_tree(const huffman_tree&);
	huffman_tree& operator = (const huffman_tree&);

protected:
	struct node {
		node *l, *r;
		bool leaf;
		byte key;
		size_t freq;
	} *root_, *cursor_;

public:
	// catch (cr::exception<huffman_tree>)
	enum errcode {
		load_damaged_stream,
		grow_damaged_stream,
		step_damaged_stream 
	};

	typedef bitstring<CODE_SPACE> bitstring;
	typedef bitstring code_table[CODE_SPACE];
	typedef size_t frequencies[CODE_SPACE];

public:
	huffman_tree()
		: root_(NULL)
		, cursor_(NULL)
	{}

	~huffman_tree()
	{
		clear(root_);
	}

public:
	void build(const frequencies& freqs)
	{
		// create forest from freq table
		TRACE(_T("****** pre-build huffman forest with frequencies ******\n"));
		std::vector<node*> trees;
		for (int i=0; i<CODE_SPACE; ++i) {
#define _SKIP_NO_FREQ_
#ifdef _SKIP_NO_FREQ_
			if (freqs[i] > 0) {
				// unless defined _SKIP_NO_FREQ_
				// all bytes have code, although freq is 0
				// for the code may used on other stream
#endif //_SKIP_NO_FREQ_
				node* tree = (node*)malloc(sizeof(node));
				tree->l = tree->r = NULL;
				tree->leaf = true;
				tree->key = (byte)i;
				tree->freq = freqs[i]; 
				trees.push_back(tree);
				TRACE(_T("init key=%0.2X, freq=%0.8X\n"), tree->key, tree->freq);
#ifdef _SKIP_NO_FREQ_
			}
#endif //_SKIP_NO_FREQ_
		}
		ASSERT(!root_);
		TRACE(_T("****** huffman forest merge to huffman tree ******\n"));
		root_ = merge(trees);
		if (!root_)
			return;
		if (root_->leaf) {
			//if only one type of byte in all bytes
			//thus the root_ is leaf
			//but root take no bit in code
			//so adjust it to left child (or right is also)
			node* tree = (node*)malloc(sizeof(node));
			tree->l = root_;
			tree->r = NULL;
			tree->leaf = false;
			//no need(tree->key = any;)
			tree->freq = root_->freq;
			root_ = tree;
		}
		cursor_ = root_;
	}

	void generate(code_table& codes)
	{
		TRACE(_T("****** huffman tree generate the codes ******\n"));
		bitstring temp;
		visit(root_, temp, codes);
	}

	void restore(const code_table& codes)
	{
		TRACE(_T("****** restore huffman tree from codes ******\n"));
		ASSERT(!root_);
		root_ = (node*)malloc(sizeof(node));
		root_->l = root_->r = NULL;
		root_->leaf = false;
		//no need(root_->key = any;)
		//no need(root_->freq = any;)
		for (short i=0; i<CODE_SPACE; ++i)
		{
			cursor_ = root_;
			const bitstring& bs = codes[i];
			int nsize = bs.length();
			if (nsize <= 0)
				continue;
			TRACE(_T("0x%02X(%d):"), i, nsize);
			DEBUGING(codes[i].dump());
			TRACE(_T("\n"));
			for (short j=0; j<nsize; ++j)
				grow(codes[i][j]);
			cursor_->key = (byte)i;
			cursor_->leaf = true;
		}
		if (!root_->r && !root_->l) {
			free(root_);
			root_ = NULL;
		}
		cursor_ = root_;
	}

protected:
	node* takeoff(std::vector<node*>& trees)
	{
		ASSERT(!trees.empty());
		std::vector<node*>::iterator it=trees.begin();
		std::vector<node*>::iterator mi = it;
		size_t tmp = (*mi)->freq;
		for (; trees.end()!=it; ++it) {
			ASSERT(*it);
			if ((*it)->freq < tmp) {
				mi = it;
				tmp = (*mi)->freq;
			}
		}
		node* tree = *mi;
		trees.erase(mi);
		ASSERT(tree);
		return tree;
	}

	node* merge(std::vector<node*>& trees)
	{
		if (trees.empty())
			return NULL;
		if (trees.size() == 1)
			return takeoff(trees);
		ASSERT(trees.size() >= 2);
		node* tree = (node*)malloc(sizeof(node));
		tree->l = takeoff(trees);
		tree->r = takeoff(trees);
		ASSERT(tree->l && tree->r);
		tree->leaf = false;
		//no need(tree->key = any;)
		tree->freq = tree->l->freq + tree->r->freq;
		TRACE(_T("%0.8X + %0.8X : %0.8X\n"),
			tree->l->freq,
			tree->r->freq,
			tree->freq);
		trees.push_back(tree);
		return merge(trees);
	}

	void visit(node* _tree, bitstring& temp, code_table& codes)
	{
		if (!_tree)
			return;
		if (_tree->leaf) {
			codes[_tree->key] = temp;
			return;
		}
		temp.append(0);
		visit(_tree->l, temp, codes);
		temp.erase(-1);
		temp.append(1);
		visit(_tree->r, temp, codes);
		temp.erase(-1);
	}

	void clear(node* _tree)
	{
		if (_tree) {
			clear(_tree->l);
			clear(_tree->r);
			free(_tree);
		}
	}

	node* copy(node* _src)
	{
		if (!_src)
			return NULL;
		node* _dst = (node*)malloc(sizeof(node));
		memcpy_s(_dst, sizeof(node), _src, sizeof(node));
		_dst->l = copy(_src->l);
		_dst->r = copy(_src->r);
		return _dst;
	}

	inline void grow(bool bi)
	{
		ASSERT(cursor_ && !cursor_->leaf);
		node* &childref = (bi ? cursor_->r : cursor_->l);
		if (!childref) {
			//childref is reference(&) to child right or left. 
			childref = (node*)malloc(sizeof(node));
			childref->l = childref->r = NULL;
			//no need(childref->freq = any;)
			//no need(childref->key = any;)
			childref->leaf = false;
		}
		if (childref->leaf)//code path conflict 
			throw exception<huffman_tree>(grow_damaged_stream);
		cursor_ = childref;
		ASSERT(!cursor_->leaf);
	}

#ifdef _DEBUG
protected:
	void probe(std::ostream& out, const node* tree, 
		int dep, int depfrom = 0, int depto = -1) const
	{
		if (!tree)
			return;
		probe(out, tree->l, dep+1, depfrom, depto);
		if (dep < depfrom)
			out << std::endl;
		else if (dep <= depto || depto < depfrom) {
			for (int i = depfrom; i < dep; i++)
				out << "          |";
			out << "+" << std::setfill('0') 
				<< std::setbase(16) << setiosflags(std::ios::uppercase) 
				<< std::setw(2) << (int)(tree->leaf ? tree->key : 0) << ":" 
				/*<< setw(5) << (int)(tree->freq)*/
				<< "-|" << std::endl;
		}
		probe(out, tree->r, dep+1, depfrom, depto);
	}
	void dump(std::ostream& out) const 
	{
		out << "c.r. test for tree.\n";
		out << "-----------------------------------\n";
		out << "             |- left child\n +key:freq-|\n             |- right child\n";
		out << "-----------------------------------\n";
		out << "depth(0 is root): " << 0 << " + n" << std::endl;
		for (int i = 0; i <= 9; i++)
			out << "+" << i << "_________";
		out << std::endl << std::endl;
		probe(out, root_, 0, 0, -1);
		out << "-----------------------------------\n";
	}
#endif //_DEBUG

};

}//namespace crhuffman

#endif//__CR_HUFFMAN_TREE_H__